"""
知识图谱构建测试
"""

# 构建数据集
import math

from torch import nn
import numpy as np

entity_list, relation_list = set(), set()
pairs = []
entity = {0, 1, 2}
relationShips = {0, 1}
pairs = [[0, 0, 1], [0, 0, 2], [1, 1, 2], [2, 0, 0], [1, 0, 0]]
n_pairs = [[0, 0, 2], [1, 0, 0], [1, 1, 0], [1, 1, 0], [2, 1, 0]]


# 构建模型
class My_Modul(nn.Module):
    def __init__(self):
        super(My_Modul, self).__init__()
        self.margin = 1  # (式5-5)中的 m
        self.e_embeding = [[0.0, 0.0, 0.0, 0.0, 0.0], [1.0, 0.0, 0.0, 4.0, 0.0], [0.0, 3.0, 0.0, 0.2, 2.0]]
        self.r_embeding = [[1.0, 2.0, 3.0, 4.0, 5.0], [0.0, 0.0, 1.1, 0.0, 0.0]]

    def forward(self, index):
        dist_correct = self.predict(pairs[index])
        dist_corrupt = self.predict(n_pairs[index])
        return self.__hinge_loss(dist_correct, dist_corrupt)

    # 返回的是差值
    def predict(self, pair):
        h = self.e_embeding[pair[0]]
        r = self.r_embeding[pair[1]]
        t = self.e_embeding[pair[2]]
        score = h + r - t
        return math.sum(score ** 2, axis=1, keepdims=True) ** 0.5

    # 损失函数，把正例差值和负例差值相加，得到损失值
    def __hinge_loss(self, dist_correct, dist_corrupt):
        a = dist_correct - dist_corrupt + self.margin
        return np.maximum(a, 0)
